import sys
file_path = sys.argv[1]



def pdb2seq(file_path):

    # file_path = sys.argv[1]
    chain_list = []
    pdb_file = open("./{}".format(file_path), 'r')
    pdb_lines = [i for i in pdb_file]
    pdb_ca = []
    pdb_AA = []
    pdb_AA_all = []
    seqs = []

    for i in pdb_lines:
        if (i[0:5].strip() == "ATOM") & (i[13:16].strip() == "CA"):
            pdb_ca.append(i)
    for i in pdb_ca:
        chain_list.append(i[21:23].strip())
    chain_list = list(set(chain_list))

    for i in chain_list:
        for j in pdb_ca:
            if j[21:23].strip() == i:
                pdb_AA.append(j[17:20])
        pdb_AA_all.append(pdb_AA)
        pdb_AA = []

    seqs = pdb_AA_all
    # print( [chain_list,seqs])
    return [chain_list,seqs]



def renumber(chain,seq,file_path):
    pdb = open(file_path, "r")
    pdb_lines = [i for i in pdb]
    pdb.close()
    atom = []
    list = []
    list_nochain = []
    #qvdiao 所以非atom
    for i in range(len(pdb_lines)):
        if pdb_lines[i][0:3]  == 'ATO':
            atom.append(pdb_lines[i])

    for i in range(len(atom)):
        if  atom[i][21] == chain:
            list.append(atom[i])
        else:
            list_nochain.append(atom[i])
    list.append("ATOM     xx  OXT GLU x  xx      31.264  43.180 -29.712  1.00 50.21           O1-  \n")
    #分割pdb_lines
    index_list = []
    # print(list)
    for i in list:
        if i[12:15].strip() == "N":
            index_list.append(1)
        else:
            index_list.append(0)
    index_list2 = [i for i in range(len(index_list)) if index_list[i] == 1]
    index_list2.append(-1)
    # print(index_list2)
    list_clip = []
    for i in range(len(index_list2)-1):
        list_clip.append(list[index_list2[i]:index_list2[i+1]])

    # print(list_clip)

    for i in range(len(list_clip)):
        for j in range(len(list_clip[i])):
            num = ""
            if len(str(i + 1)) == 1:
                num = "   " + str(i + 1)
            if len(str(i + 1)) == 2:
                num = "  " + str(i + 1)
            if len(str(i + 1)) == 3:
                num = " " + str(i + 1)
            if len(str(i + 1)) == 4:
                num = str(i)

            list_clip[i][j] = list_clip[i][j][0:22] + num + list_clip[i][j][26:-1] + list_clip[i][j][-1]

    # print(list_clip)

    merge = []
    for i in  list_clip:
        merge += i

    pdb_file_list  = list_nochain + merge
    # print(pdb_file_list)
    f = open("./renumbered.pdb","w+")
    f.writelines(pdb_file_list)
    f.close()


pdb2seq = pdb2seq(file_path)
chain_list = pdb2seq[0]
seqs = pdb2seq[1]
# print(chain_list)
if len(chain_list) > 1 :
    for i in range(len(chain_list)):
        renumber(chain_list[i],seqs[i],file_path)
        file_path = "./renumbered.pdb".format(file_path)
elif len(chain_list) == 1 :
    renumber(chain_list[0], seqs[0], file_path)




